# SQL 生成 Agent

from .base_agent import BaseAgent
from models.state import WorkflowState
# from utils.database import get_db_connection
from zai import ZhipuAiClient
import re

class SQLGenerationAgent(BaseAgent):
    def __init__(self):
        self.client = ZhipuAiClient(api_key="db2e03cca9654d4dba9b91b123548374.TftHaYZq4ydQPDjP")
    
    def process(self, state: WorkflowState) -> WorkflowState:
        """
        根据意图和实体生成SQL语句
        """
        user_input = state.user_input
        intent = state.intent
        entities = state.entities or {}
        schema_info = state.schema_info or {}
        
        print(f"SQL生成Agent开始工作...")
        print(f"意图类型: {intent}")
        print(f"用户输入: {user_input}")
        
        # 根据不同的意图类型生成相应的SQL
        if intent == "Imperativequery":
            sql = self._generate_imperative_sql(user_input, entities, schema_info)
        elif intent == "ExploratoryAnalysis":
            sql = self._generate_exploratory_sql(user_input, entities, schema_info)
        elif intent == "Dataextraction":
            sql = self._generate_data_extraction_sql(user_input, entities, schema_info)
        else:
            sql = None
        
        state.generated_sql = sql
        return state
    
    def _generate_imperative_sql(self, user_input: str, entities: dict, schema_info: dict) -> str:
        """
        生成指令式查询的SQL
        - 事实查询、状态查询、简单统计
        """
        prompt = """你是一个SQL专家，请根据用户问题生成相应的SQL查询语句。

【指令式查询特点】
- 获取具体数值或简单统计
- 返回单个值或简单列表
- 包含聚合函数：COUNT, SUM, AVG, MAX, MIN等

【数据库表结构】
假设有以下表：
- sales表: id, product_name, sale_date, amount, quantity, region
- users表: id, username, register_date, status
- products表: id, product_name, category, price, stock

【要求】
1. 只返回SQL语句，不要解释
2. 使用合适的聚合函数
3. 包含必要的WHERE条件
4. 使用标准SQL语法

用户问题："{user_input}"

SQL语句："""
        
        try:
            response = self.client.chat.completions.create(
                model="glm-4",
                messages=[{"role": "user", "content": prompt.format(user_input=user_input)}],
                temperature=0.1,
                max_tokens=200
            )
            
            sql = self._extract_sql_from_response(response)
            return self._validate_and_clean_sql(sql)
            
        except Exception as e:
            print(f"生成指令式SQL出错: {e}")
            return self._generate_fallback_sql(user_input, "imperative")
    
    def _generate_exploratory_sql(self, user_input: str, entities: dict, schema_info: dict) -> str:
        """
        生成探索式分析的SQL
        - 数据探索、模式发现、趋势分析
        """
        prompt = """你是一个SQL专家，请根据用户问题生成探索式分析的SQL查询语句。

【探索式分析特点】
- 分析数据趋势和模式
- 包含GROUP BY分组
- 包含ORDER BY排序
- 可能包含JOIN操作
- 用于发现洞察而非获取单一数值

【数据库表结构】
假设有以下表：
- sales表: id, product_name, sale_date, amount, quantity, region, category
- users表: id, username, register_date, status, region
- products表: id, product_name, category, price, stock

【要求】
1. 只返回SQL语句，不要解释
2. 使用GROUP BY进行分组分析
3. 使用ORDER BY进行排序
4. 包含有意义的列选择
5. 使用标准SQL语法

用户问题："{user_input}"

SQL语句："""
        
        try:
            response = self.client.chat.completions.create(
                model="glm-4",
                messages=[{"role": "user", "content": prompt.format(user_input=user_input)}],
                temperature=0.1,
                max_tokens=200
            )
            
            sql = self._extract_sql_from_response(response)
            return self._validate_and_clean_sql(sql)
            
        except Exception as e:
            print(f"生成探索式SQL出错: {e}")
            return self._generate_fallback_sql(user_input, "exploratory")
    
    def _generate_data_extraction_sql(self, user_input: str, entities: dict, schema_info: dict) -> str:
        """
        生成数据提取的SQL
        - 记录查询、数据导出、报表生成
        """
        prompt = """你是一个SQL专家，请根据用户问题生成数据提取的SQL查询语句。

【数据提取特点】
- 获取完整的记录列表
- 可能包含所有列或特定列
- 用于导出或生成报表
- 可能包含分页(LIMIT)或条件过滤

【数据库表结构】
假设有以下表：
- sales表: id, product_name, sale_date, amount, quantity, region, customer_id
- users表: id, username, email, register_date, status, region
- products表: id, product_name, category, price, stock, description

【要求】
1. 只返回SQL语句，不要解释
2. 选择适当的列（避免SELECT *）
3. 包含必要的WHERE条件
4. 考虑数据量，合理使用LIMIT
5. 使用标准SQL语法

用户问题："{user_input}"

SQL语句："""
        
        try:
            response = self.client.chat.completions.create(
                model="glm-4",
                messages=[{"role": "user", "content": prompt.format(user_input=user_input)}],
                temperature=0.1,
                max_tokens=200
            )
            
            sql = self._extract_sql_from_response(response)
            return self._validate_and_clean_sql(sql)
            
        except Exception as e:
            print(f"生成数据提取SQL出错: {e}")
            return self._generate_fallback_sql(user_input, "extraction")
    
    def _extract_sql_from_response(self, response) -> str:
        """从大模型响应中提取SQL语句"""
        if hasattr(response, 'choices') and len(response.choices) > 0:
            if hasattr(response.choices[0], 'message'):
                content = response.choices[0].message.content.strip()
                # 提取SQL语句（去除可能的代码块标记）
                sql_match = re.search(r'```sql\s*(.*?)\s*```', content, re.DOTALL)
                if sql_match:
                    return sql_match.group(1).strip()
                else:
                    # 如果没有代码块，直接返回内容
                    return content
        return ""
    
    def _validate_and_clean_sql(self, sql: str) -> str:
        """验证和清理SQL语句"""
        if not sql:
            return "SELECT 1"  # 默认返回有效SQL
        
        # 移除多余的空白字符
        sql = re.sub(r'\s+', ' ', sql).strip()
        
        # 确保以SELECT开头（安全限制）
        if not sql.upper().startswith('SELECT'):
            sql = "SELECT 1"
            
        return sql
    
    def _generate_fallback_sql(self, user_input: str, intent_type: str) -> str:
        """生成备用SQL（当大模型调用失败时）"""
        fallback_sqls = {
            "imperative": "SELECT COUNT(*) as total_count FROM sales WHERE sale_date >= DATE_SUB(NOW(), INTERVAL 1 YEAR)",
            "exploratory": "SELECT product_name, SUM(amount) as total_sales FROM sales GROUP BY product_name ORDER BY total_sales DESC LIMIT 10",
            "extraction": "SELECT * FROM users WHERE register_date >= DATE_SUB(NOW(), INTERVAL 3 MONTH) LIMIT 100"
        }
        return fallback_sqls.get(intent_type, "SELECT 1")